import tensorflow as tf
import os
from tensorflow.python.framework.graph_util import convert_variables_to_constants

model_root_path = ''


def create_model():
    test_name = 'tf.math.argmax'
    tf.compat.v1.disable_eager_execution()
    x = tf.compat.v1.placeholder(tf.int16, shape=(14, 5, 3, 4, 12), name="Input")
    y = tf.math.argmax(input=x, axis=0, name="output")
    output = tf.convert_to_tensor(y, name="Output")
    with tf.compat.v1.Session() as sess:
        sess.run(tf.compat.v1.global_variables_initializer())
        graph = convert_variables_to_constants(sess, sess.graph_def, ["output"])
        tf.io.write_graph(graph, '.', os.path.join(model_root_path, test_name + '.pb'), as_text=False)
        print('Create Model Successful.')
        print('Path: ', os.path.join(model_root_path, test_name + '.pb'))

    tf.compat.v1.reset_default_graph()


if __name__ == "__main__":
    create_model()